import argparse
import pandas as pd
import json
from rouge_score import rouge_scorer, tokenizers
import torch
from transformers import BertTokenizer, BertModel
from sklearn.metrics.pairwise import cosine_similarity
import jieba

# 自定义分词器
class JiebaTokenizer(tokenizers.Tokenizer):
    def tokenize(self, text):
        return list(jieba.cut(text))

# 初始化 ROUGE 计算器，使用自定义分词器
scorer = rouge_scorer.RougeScorer(['rougeL'], tokenizer=JiebaTokenizer(), use_stemmer=False)

def calculate_rouge_l(text1, text2):
    """
    计算 ROUGE-L 分数
    """
    # 计算 ROUGE-L 分数
    scores = scorer.score(text1, text2)
    print(f"ROUGE-L 分数: {scores['rougeL']}, text1= {text1}, text2= {text2}")  # 打印调试信息
    return scores['rougeL'].fmeasure

def calculate_bert_similarity(text1, text2, tokenizer, model, device):
    """
    计算 BERT 语义相似度
    """
    # 分词并编码
    inputs1 = tokenizer(text1, return_tensors='pt', padding=True, truncation=True).to(device)
    inputs2 = tokenizer(text2, return_tensors='pt', padding=True, truncation=True).to(device)

    # 获取句向量
    with torch.no_grad():
        outputs1 = model(**inputs1)
        outputs2 = model(**inputs2)

    # 取 [CLS] 位置的向量作为句向量
    sentence_embedding1 = outputs1.last_hidden_state[:, 0, :].cpu().numpy()
    sentence_embedding2 = outputs2.last_hidden_state[:, 0, :].cpu().numpy()

    # 计算余弦相似度
    similarity = cosine_similarity(sentence_embedding1, sentence_embedding2)[0][0]
    return similarity

def calculate_opsim(text1, text2, tokenizer, model, device, alpha=0.5):
    """
    计算 OpSim 指标
    :param text1: 文本1
    :param text2: 文本2
    :param tokenizer: BERT 分词器
    :param model: BERT 模型
    :param device: 设备（CPU 或 GPU）
    :param alpha: ROUGE-L 的权重
    :return: OpSim 值
    """
    rouge_l = calculate_rouge_l(text1, text2)
    bert_similarity = calculate_bert_similarity(text1, text2, tokenizer, model, device)
    print(f"rouge_l: {rouge_l}, bert_similarity: {bert_similarity}")
    opsim = alpha * rouge_l + (1 - alpha) * bert_similarity
    return opsim

def load_data(file_path, file_type):
    """
    根据文件类型加载数据
    """
    if file_type == "json":
        with open(file_path, "r", encoding="utf-8") as f:
            data = json.load(f)
        if not isinstance(data, list) or not all("opinion_origin" in item and "opinion_new" in item for item in data):
            raise ValueError("JSON 文件必须是一个包含 'opinion_origin' 和 'opinion_new' 键的列表")
        return data
    elif file_type == "excel":
        df = pd.read_excel(file_path)
        if "opinion_origin" not in df.columns or "opinion_new" not in df.columns:
            raise ValueError("Excel 文件中必须包含 'opinion_origin' 和 'opinion_new' 列")
        return df.to_dict(orient="records")
    else:
        raise ValueError("不支持的文件类型，请指定 'json' 或 'excel'")

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="计算 OpSim 指标")
    parser.add_argument("--type", type=str, required=True, choices=["json", "excel"], help="文件类型（json 或 excel）")
    parser.add_argument("--test_file", type=str, required=True, help="测试文件路径")
    parser.add_argument("--model_path", type=str, required=True, help="BERT 模型路径")
    args = parser.parse_args()

    # 检查 GPU 是否可用
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # 加载 BERT 模型和分词器
    tokenizer = BertTokenizer.from_pretrained(args.model_path)
    model = BertModel.from_pretrained(args.model_path).to(device)  # 将模型移动到 GPU

    # 加载数据
    data = load_data(args.test_file, args.type)

    # 计算每对文本的 OpSim 指标
    opsim_scores = []
    for item in data:
        text1 = item["opinion_origin"]
        text2 = item["opinion_new"]
        opsim = calculate_opsim(text1, text2, tokenizer, model, device, alpha=0.5)
        opsim_scores.append(opsim)

    # 计算平均值
    average_opsim = sum(opsim_scores) / len(opsim_scores)
    print(f"所有项的 OpSim 平均值: {average_opsim:.4f}")

if __name__ == "__main__":
    main()